import numpy as np
import torch
from tqdm import tqdm
import torch.nn.functional as F


def mean_iou(input, target, classes = 10):
    """  compute the value of mean iou
    :param input:  2d array, int, prediction
    :param target: 2d array, int, ground truth
    :param classes: int, the number of class
    :return:
        miou: float, the value of miou
    """
    miou = []
    for i in range(classes):
        intersection = np.logical_and(target == i, input == i)
        # print(intersection.any())
        union = np.logical_or(target == i, input == i)
        temp = np.sum(intersection) / np.sum(union)
        miou += [temp]
    return  np.array(miou)[np.newaxis,...]
def val_fun(model,data_loader):
    all_iou = np.zeros((1,10))
    model.eval()
    with torch.no_grad():
        for batch_image,batch_label in tqdm(data_loader):
            batch_pred = model(batch_image.float().cuda())
            batch_pred = torch.argmax(F.softmax(batch_pred, dim=1), dim=1)
            batch_pred = batch_pred.cpu().numpy()
            for batch_idx in range(batch_pred.shape[0]):
                iou = mean_iou(batch_pred[batch_idx],batch_label[batch_idx].cpu().numpy(),classes=10)
                # print(iou)
                all_iou = np.concatenate((all_iou,iou),axis=0)
    return np.nanmean(all_iou,axis=0)

if __name__ == '__main__':
    import torch
    from torch.utils.data import DataLoader
    from model import model_unet_res50
    from dataset import MyDataset

    model = model_unet_res50.cuda()
    model.load_state_dict(torch.load('../model_save/new/unet_res50_baseline/best.pth'))
    DataPath = '../suichang_round1_train_210120/suichang_round1_train_210120'
    TrainJsonDir = 'label_file/train.json'
    ValJsonDir = 'label_file/train.json'
    val_dataset = MyDataset(data_path=DataPath, name_list_dir=ValJsonDir, transform=None)
    val_dataloader = DataLoader(dataset=val_dataset, batch_size=2, shuffle=False, num_workers=0)
    viou = val_fun(model=model,data_loader=val_dataloader)
    print('\t'.join(viou.astype(str)))
    print(np.mean(viou))